"""
对训练集进行2D裁剪，每张ct抽取10张作为有标注数据，其他的作为无标注数据
并将标注和无标注的文件id用txt文件保存下来，2D-cut后的数据放在一个文件夹下
"""
import os
import random
import sys
sys.path.append(os.path.split(sys.path[0])[0])
import SimpleITK as sitk
from multiprocessing.dummy import Pool
from config import parameter as para
from data_prepare.dataProcess import processLiverOrTumorMask
import numpy as np
from utilities.StatisticTime import getTime
# for file in tqdm(os.listdir(para.vol_path)[:para.train_test_split]):
def process(file):
    # 将CT和标签入读内存
    volume = sitk.ReadImage(os.path.join(para.niivol_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    #print('处理前的volume shape:',volume_array.shape,end = '   ')
    seg = sitk.ReadImage(os.path.join(para.niiseg_path, file.replace('volume', 'segmentation')), sitk.sitkUInt8)
    seg_array = sitk.GetArrayFromImage(seg)
    seg_file_name=file.replace('volume', 'segmentation')
    ct_file_name=file
    # seg_array=processLiverOrTumorMask(para.Liver_or_Tumor,seg_array)
    # 将灰度值在阈值之外的截断掉
    volume_array[volume_array > para.upper] = para.upper
    volume_array[volume_array < para.lower] = para.lower

    tumor_indexs=get_sampling(seg_array,filename = ct_file_name)#选取的所有下标
    print(ct_file_name,len(tumor_indexs),volume_array.shape[0],len(tumor_indexs)/volume_array.shape[0])

    return len(tumor_indexs)

def get_sampling(seg_array,filename):

    tumor_indexs = list(np.where(np.any(seg_array == 2, axis = (1, 2)))[0])  # 肿瘤区域

    return tumor_indexs

@getTime
def main():
    # 需要训练集的file_list
    with open(para.train_nii_id_path, 'r') as f:
        labeled_ids = f.read().splitlines()
    file_list = [os.path.basename(file.split()[0]) for file in labeled_ids]
    # 多线程
    pool = Pool(1)
    result = pool.map(process, file_list)
    pool.close()
    pool.join()
    print(result)
if __name__ == '__main__':
    #
    # if os.path.exists(para.cut_save_path):
    #     shutil.rmtree(para.cut_save_path)
   # main()
    l=[36, 51, 22, 38, 48, 87, 0, 45, 199, 68, 0, 45, 14, 133, 27, 64, 235, 9, 131, 15, 16, 211, 18, 4, 7, 35, 53, 85, 5, 102, 18, 11, 48, 10, 136, 119, 34, 7, 76, 0, 22, 12, 120, 11, 199, 9, 37, 17, 0, 0, 0, 4, 10, 11, 18, 49, 30, 244, 32, 11, 122, 190, 0, 7, 30, 66, 53, 18, 15, 192, 64, 81, 4, 10, 0, 27, 78, 29, 36, 20, 93, 141, 60, 5, 226, 130, 80, 54, 22, 83, 74, 7, 72, 7, 21, 92, 10, 0, 9, 0]
    print(np.array(l).mean())

